from itertools import combinations
from collections import defaultdict
class Apriori:
    def __init__(self, min_support = 0.5, min_confidence = 0.5):
        self.min_support = min_support
        self.min_confidence = min_confidence
        self.freq_itemsets = None
        self.support_data = {}
        self.transactions = None
    def fit(self, transactions):
        """ 执行Apriori算法来发现频繁项集 """
        # 将交易中的每个商品转换为集合
        self.transactions = [set(transaction) for transaction in transactions]
        itemset = set([item for transaction in self.transactions for item in transaction])
        # 初始化self.freq_itemsets为一个空列表
        self.freq_itemsets = []
        self.find_frequent_itemsets(itemset, 1)
    def find_frequent_itemsets(self, itemset, k):
        """ 找到频繁k项集 """
        if k > 1:
            itemset = set([frozenset(i) for i in combinations(itemset, k)])
            itemset = self.prune(itemset)
        # 当k=1时，确保item是单元素集合
        if k == 1:
            freq_itemset = defaultdict(int)
            for transaction in self.transactions:
                for item in itemset:
                    if {item}.issubset(transaction):  # 注意这里的变化
                        freq_itemset[frozenset([item])] += 1
        else:
            freq_itemset = defaultdict(int)
            for transaction in self.transactions:
                for item in itemset:
                    if item.issubset(transaction):
                        freq_itemset[item] += 1
        num_items = float(len(self.transactions))
        retlist = []
        for key, value in freq_itemset.items():
            support = value / num_items
            if support >= self.min_support:
                retlist.insert(0, key)
            self.support_data[key] = support
        # 将找到的频繁项集添加到self.freq_itemsets中
        self.freq_itemsets.extend(retlist)
        # 确保retlist不为空再进行递归调用
        if len(retlist) > 0:
            self.find_frequent_itemsets(frozenset.union(*retlist), k + 1)
    def prune(self, itemset):
        """ 剪枝过程 """
        # 如果itemset为空，直接返回空列表
        if not itemset:
            return []
        retlist = []
        k = len(list(itemset)[0])
        for item in itemset:
            subsets = combinations(item, k - 1)
            for subset in subsets:
                if not frozenset(subset) in self.freq_itemsets:
                    break
            else:
                retlist.append(item)
        return retlist
    def generate_rules(self):
        """ 生成关联规则 """
        rules = []
        for itemset in self.freq_itemsets:
            if len(itemset) > 1:
                for i in range(1, len(itemset)):
                    for subset in combinations(itemset, i):
                        subset = frozenset(subset)
                        remain = itemset - subset
                        confidence = self.support_data[itemset] / self.support_data[subset]
                        if confidence >= self.min_confidence:
                            rules.append((subset, remain, confidence))
        return rules
# 示例数据集
transactions = [
    ['牛奶', '面包', '啤酒'],
    ['面包', '咖啡', '牛奶'],
    ['面包', '咖啡', '牛奶'],
    ['牛奶', '咖啡'],
    ['面包', '咖啡', '牛奶'],
    ['面包', '咖啡', '牛奶'],
    ['牛奶', '面包', '啤酒'],
    ['面包', '牛奶'],
    ['面包', '咖啡', '牛奶'],
    ['面包', '咖啡', '牛奶']
]
# 训练模型
ap = Apriori(min_support = 0.5, min_confidence = 0.5)
ap.fit(transactions)
# 输出频繁项集
print("频繁项集:")
for itemset in ap.freq_itemsets:
    print(f"{itemset}: {ap.support_data[itemset]}")
# 输出关联规则
print("\n关联规则:")
rules = ap.generate_rules()
for rule in rules:
    print(f"{rule[0]} -> {rule[1]} (confidence: {rule[2]:.2f})")